"""
音频流处理和缓冲管理
支持实时音频流传输和缓冲
"""

import asyncio
import io
import time
from typing import Optional, Dict, Any, Callable, List, Tuple
from dataclasses import dataclass, field
from collections import deque
import numpy as np
import soundfile as sf
from enum import Enum


class StreamStatus(Enum):
    """流状态"""
    IDLE = "idle"
    BUFFERING = "buffering"
    PLAYING = "playing"
    PAUSED = "paused"
    STOPPED = "stopped"
    ERROR = "error"


@dataclass
class StreamConfig:
    """流配置"""
    buffer_size: int = 8192  # 缓冲区大小（字节）
    max_buffers: int = 10     # 最大缓冲数量
    sample_rate: int = 22050   # 采样率
    channels: int = 1         # 声道数
    format: str = 'wav'       # 音频格式
    chunk_duration: float = 0.1  # 每个块时长（秒）
    preload_chunks: int = 3   # 预加载块数量
    auto_reconnect: bool = True  # 自动重连
    reconnect_delay: float = 1.0  # 重连延迟


@dataclass
class AudioChunk:
    """音频块"""
    data: bytes
    timestamp: float
    duration: float
    sequence: int
    metadata: Dict[str, Any] = field(default_factory=dict)


class AudioBuffer:
    """音频缓冲区"""
    
    def __init__(self, config: StreamConfig):
        self.config = config
        self._buffer = deque(maxlen=config.max_buffers)
        self._lock = asyncio.Lock()
        self._total_size = 0
        self._sequence_counter = 0
    
    async def add_chunk(self, data: bytes, duration: float) -> bool:
        """添加音频块"""
        async with self._lock:
            if len(data) > self.config.buffer_size:
                return False
            
            chunk = AudioChunk(
                data=data,
                timestamp=time.time(),
                duration=duration,
                sequence=self._sequence_counter
            )
            
            self._sequence_counter += 1
            self._buffer.append(chunk)
            self._total_size += len(data)
            
            return True
    
    async def get_chunk(self) -> Optional[AudioChunk]:
        """获取音频块"""
        async with self._lock:
            if not self._buffer:
                return None
            
            chunk = self._buffer.popleft()
            self._total_size -= len(chunk.data)
            return chunk
    
    async def get_chunks(self, count: int) -> List[AudioChunk]:
        """获取多个音频块"""
        async with self._lock:
            chunks = []
            for _ in range(min(count, len(self._buffer))):
                chunk = self._buffer.popleft()
                self._total_size -= len(chunk.data)
                chunks.append(chunk)
            return chunks
    
    def is_empty(self) -> bool:
        """检查缓冲区是否为空"""
        return len(self._buffer) == 0
    
    def is_full(self) -> bool:
        """检查缓冲区是否已满"""
        return len(self._buffer) >= self.config.max_buffers
    
    def size(self) -> int:
        """获取缓冲区大小"""
        return len(self._buffer)
    
    def get_buffer_duration(self) -> float:
        """获取缓冲区总时长"""
        return sum(chunk.duration for chunk in self._buffer)
    
    def clear(self):
        """清空缓冲区"""
        self._buffer.clear()
        self._total_size = 0


class StreamManager:
    """音频流管理器"""
    
    def __init__(self, config: StreamConfig):
        self.config = config
        self.status = StreamStatus.IDLE
        self.buffer = AudioBuffer(config)
        self._lock = asyncio.Lock()
        self._callbacks: Dict[str, List[Callable]] = {
            'on_chunk': [],
            'on_buffer_update': [],
            'on_status_change': [],
            'on_error': []
        }
        self._streaming_task: Optional[asyncio.Task] = None
        self._metadata: Dict[str, Any] = {}
    
    async def start_streaming(self, source: Callable) -> bool:
        """开始流传输"""
        async with self._lock:
            if self.status in [StreamStatus.PLAYING, StreamStatus.BUFFERING]:
                return False
            
            self.status = StreamStatus.BUFFERING
            await self._notify_status_change()
            
            try:
                # 创建流传输任务
                self._streaming_task = asyncio.create_task(
                    self._streaming_loop(source)
                )
                return True
            except Exception as e:
                await self._handle_error(f"Failed to start streaming: {str(e)}")
                return False
    
    async def stop_streaming(self):
        """停止流传输"""
        async with self._lock:
            if self._streaming_task:
                self._streaming_task.cancel()
                try:
                    await self._streaming_task
                except asyncio.CancelledError:
                    pass
                self._streaming_task = None
            
            self.status = StreamStatus.STOPPED
            self.buffer.clear()
            await self._notify_status_change()
    
    async def pause_streaming(self):
        """暂停流传输"""
        async with self._lock:
            if self.status == StreamStatus.PLAYING:
                self.status = StreamStatus.PAUSED
                await self._notify_status_change()
    
    async def resume_streaming(self):
        """恢复流传输"""
        async with self._lock:
            if self.status == StreamStatus.PAUSED:
                self.status = StreamStatus.PLAYING
                await self._notify_status_change()
    
    async def _streaming_loop(self, source: Callable):
        """流传输循环"""
        reconnect_count = 0
        max_reconnect_attempts = 5
        
        while True:
            try:
                async for chunk_data, duration in self._generate_chunks(source):
                    if self.status == StreamStatus.PAUSED:
                        await asyncio.sleep(0.1)
                        continue
                    
                    if self.status == StreamStatus.STOPPED:
                        break
                    
                    # 检查缓冲区是否已满
                    if self.buffer.is_full():
                        await asyncio.sleep(0.01)
                        continue
                    
                    # 添加到缓冲区
                    success = await self.buffer.add_chunk(chunk_data, duration)
                    if success:
                        await self._notify_chunk_added(chunk_data, duration)
                    
                    # 更新状态
                    if self.status == StreamStatus.BUFFERING and \
                       self.buffer.size() >= self.config.preload_chunks:
                        self.status = StreamStatus.PLAYING
                        await self._notify_status_change()
                    
                    reconnect_count = 0  # 重置重连计数
                
                # 流正常结束
                await self.stop_streaming()
                break
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                if self.config.auto_reconnect and reconnect_count < max_reconnect_attempts:
                    reconnect_count += 1
                    await self._handle_error(f"Stream error, reconnecting ({reconnect_count}/{max_reconnect_attempts}): {str(e)}")
                    await asyncio.sleep(self.config.reconnect_delay)
                    continue
                else:
                    await self._handle_error(f"Stream failed: {str(e)}")
                    break
    
    async def _generate_chunks(self, source: Callable):
        """生成音频块"""
        try:
            if asyncio.iscoroutinefunction(source):
                audio_source = await source()
            else:
                audio_source = source()
            
            if isinstance(audio_source, bytes):
                # 处理字节数据
                await self._generate_chunks_from_bytes(audio_source)
            elif isinstance(audio_source, tuple) and len(audio_source) == 2:
                # 处理(numpy数组, 采样率)
                data, sample_rate = audio_source
                await self._generate_chunks_from_array(data, sample_rate)
            else:
                raise ValueError("Unsupported source type")
                
        except Exception as e:
            raise RuntimeError(f"Failed to generate chunks: {str(e)}")
    
    async def _generate_chunks_from_bytes(self, data: bytes):
        """从字节数据生成块"""
        chunk_size = int(self.config.sample_rate * self.config.chunk_duration * 2)  # 16-bit音频
        
        for i in range(0, len(data), chunk_size):
            chunk_data = data[i:i + chunk_size]
            if chunk_data:
                duration = len(chunk_data) / (self.config.sample_rate * 2)
                yield chunk_data, duration
                await asyncio.sleep(0)  # 让出控制权
    
    async def _generate_chunks_from_array(self, data: np.ndarray, sample_rate: int):
        """从numpy数组生成块"""
        samples_per_chunk = int(sample_rate * self.config.chunk_duration)
        
        for i in range(0, len(data), samples_per_chunk):
            chunk_data = data[i:i + samples_per_chunk]
            
            # 转换为字节数据
            with io.BytesIO() as buffer:
                sf.write(buffer, chunk_data, sample_rate, format='wav')
                chunk_bytes = buffer.getvalue()
            
            duration = len(chunk_data) / sample_rate
            yield chunk_bytes, duration
            await asyncio.sleep(0)
    
    async def read_chunk(self) -> Optional[AudioChunk]:
        """读取一个音频块"""
        return await self.buffer.get_chunk()
    
    async def read_chunks(self, count: int) -> List[AudioChunk]:
        """读取多个音频块"""
        return await self.buffer.get_chunks(count)
    
    def get_buffer_info(self) -> Dict[str, Any]:
        """获取缓冲区信息"""
        return {
            'size': self.buffer.size(),
            'max_size': self.config.max_buffers,
            'duration': self.buffer.get_buffer_duration(),
            'is_full': self.buffer.is_full(),
            'is_empty': self.buffer.is_empty()
        }
    
    def get_stream_info(self) -> Dict[str, Any]:
        """获取流信息"""
        return {
            'status': self.status.value,
            'config': {
                'sample_rate': self.config.sample_rate,
                'channels': self.config.channels,
                'format': self.config.format,
                'chunk_duration': self.config.chunk_duration
            },
            'buffer': self.get_buffer_info(),
            'metadata': self._metadata
        }
    
    def add_callback(self, event: str, callback: Callable):
        """添加事件回调"""
        if event in self._callbacks:
            self._callbacks[event].append(callback)
    
    def remove_callback(self, event: str, callback: Callable):
        """移除事件回调"""
        if event in self._callbacks and callback in self._callbacks[event]:
            self._callbacks[event].remove(callback)
    
    async def _notify_chunk_added(self, data: bytes, duration: float):
        """通知音频块添加事件"""
        for callback in self._callbacks['on_chunk']:
            try:
                if asyncio.iscoroutinefunction(callback):
                    await callback(data, duration)
                else:
                    callback(data, duration)
            except Exception as e:
                print(f"Callback error: {e}")
    
    async def _notify_buffer_update(self):
        """通知缓冲区更新事件"""
        for callback in self._callbacks['on_buffer_update']:
            try:
                if asyncio.iscoroutinefunction(callback):
                    await callback(self.get_buffer_info())
                else:
                    callback(self.get_buffer_info())
            except Exception as e:
                print(f"Callback error: {e}")
    
    async def _notify_status_change(self):
        """通知状态变化事件"""
        for callback in self._callbacks['on_status_change']:
            try:
                if asyncio.iscoroutinefunction(callback):
                    await callback(self.status)
                else:
                    callback(self.status)
            except Exception as e:
                print(f"Callback error: {e}")
    
    async def _handle_error(self, error_message: str):
        """处理错误"""
        self.status = StreamStatus.ERROR
        for callback in self._callbacks['on_error']:
            try:
                if asyncio.iscoroutinefunction(callback):
                    await callback(error_message)
                else:
                    callback(error_message)
            except Exception as e:
                print(f"Error callback error: {e}")
    
    def set_metadata(self, key: str, value: Any):
        """设置元数据"""
        self._metadata[key] = value
    
    def get_metadata(self, key: str) -> Any:
        """获取元数据"""
        return self._metadata.get(key)
    
    async def cleanup(self):
        """清理资源"""
        await self.stop_streaming()
        self.buffer.clear()
        self._callbacks.clear()
